import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QLabel, QPushButton,
    QVBoxLayout, QWidget, QFileDialog, QMessageBox, QHBoxLayout
)
from PyQt5.QtGui import QImage, QPixmap, QPainter, QColor, QPalette
from PyQt5.QtCore import Qt, QPoint


class RoundedButton(QPushButton):
    """自定义圆角按钮"""

    def __init__(self, text, parent=None):
        super().__init__(text, parent)
        self.setStyleSheet("""
            QPushButton {
                background-color: #4a90e2;
                color: white;
                border-radius: 10px;
                padding: 16px 32px;
                font-size: 20px;
                font-weight: 500;
                border: none;
            }
            QPushButton:hover {
                background-color: #0961D1;
            }
            QPushButton:pressed {
                background-color: #2a70c2;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
        """)
        self.setMinimumHeight(35)


class RoundedImageLabel(QLabel):
    """自定义圆角图像显示标签"""

    def __init__(self, text, parent=None):
        super().__init__(text, parent)
        self.setAlignment(Qt.AlignCenter)
        self.setStyleSheet("""
            QLabel {
                background-color: #f0f0f0;
                border-radius: 10px;
                border: 2px dashed #cccccc;
                color: #605353;
                font-size: 20px;
            }
        """)


class GrabCutApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("图片抠图工具")
        self.setMinimumSize(800, 600)

        # 设置窗口样式
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f8f9fa;
            }
        """)

        # 图像和标记
        self.img = None
        self.img_show = None
        self.img_gc = None
        self.mask = None
        self.drawing = False
        self.last_point = None
        self.firt_choose = True
        self.lb_down = False
        self.rb_down = False
        self.result_ready = False

        # 创建主布局
        self.init_ui()

    def init_ui(self):
        # 图像显示区域
        self.image_label = RoundedImageLabel("请点击上传图片")
        self.image_label.setFixedSize(700, 450)

        # 按钮
        btn_upload = RoundedButton("上传图片")
        btn_save = RoundedButton("保存结果")
        btn_close = RoundedButton("关闭")

        # 按钮功能绑定
        btn_upload.clicked.connect(self.upload_image)
        btn_save.clicked.connect(self.save_result)
        btn_close.clicked.connect(self.close)

        # 按钮布局
        hbox = QHBoxLayout()
        hbox.addStretch(1)
        hbox.addWidget(btn_upload)
        hbox.addSpacing(20)
        hbox.addWidget(btn_save)
        hbox.addSpacing(20)
        hbox.addWidget(btn_close)
        hbox.addStretch(1)

        # 主布局
        vbox = QVBoxLayout()
        vbox.addStretch(2)
        vbox.addWidget(self.image_label, alignment=Qt.AlignCenter)
        vbox.addStretch(1)
        vbox.addLayout(hbox)
        vbox.addStretch(2)

        # 设置中心部件
        container = QWidget()
        container.setLayout(vbox)
        self.setCentralWidget(container)

        # 事件代理
        self.image_label.mousePressEvent = self.mouse_press
        self.image_label.mouseMoveEvent = self.mouse_move
        self.image_label.mouseReleaseEvent = self.mouse_release

    def upload_image(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "Images (*.png *.jpg *.jpeg *.bmp)"
        )
        if path:
            self.img = cv2.imread(path)
            if self.img is None:
                QMessageBox.warning(self, "错误", "无法加载图片！")
                return
            self.img_show = self.img.copy()
            self.img_gc = self.img.copy()
            self.mask = np.full(self.img.shape[:2], 2, dtype=np.uint8)
            self.firt_choose = True
            self.result_ready = False
            self.update_display()
            # QMessageBox.information(self, "提示", f"已加载图片: {path}")

    def save_result(self):
        if not self.result_ready:
            QMessageBox.warning(self, "错误", "请先完成分离操作！")
            return
        alpha = np.where((self.mask == 0) | (self.mask == 2), 0, 255).astype(np.uint8)
        b, g, r = cv2.split(self.img_gc)
        bgra = cv2.merge([b, g, r, alpha])
        save_path, _ = QFileDialog.getSaveFileName(self, "保存图片", "target.png", "PNG Files (*.png)")
        if save_path:
            cv2.imwrite(save_path, bgra)
            QMessageBox.information(self, "提示", f"保存成功: {save_path}")

    def mouse_press(self, event):
        if self.img is None:
            return
        x, y = event.pos().x(), event.pos().y()
        x, y = self.map_to_image(x, y)
        if x < 0 or y < 0:
            return

        self.drawing = True
        self.last_point = (x, y)

        if event.button() == Qt.LeftButton:
            self.lb_down = True
        elif event.button() == Qt.RightButton:
            if self.firt_choose:
                QMessageBox.warning(self, "提示", "请先标记前景！")
                return
            self.rb_down = True

    def mouse_move(self, event):
        if not self.drawing or self.img is None:
            return
        x, y = event.pos().x(), event.pos().y()
        x, y = self.map_to_image(x, y)
        if x < 0 or y < 0:
            return

        if self.lb_down:
            # 红色标记前景
            cv2.line(self.img_show, self.last_point, (x, y), (0, 0, 255), 10)
            cv2.rectangle(self.mask, self.last_point, (x, y), 1, -1, 4)
        elif self.rb_down:
            # 蓝色标记背景
            cv2.line(self.img_show, self.last_point, (x, y), (255, 0, 0), 10)
            cv2.rectangle(self.mask, self.last_point, (x, y), 0, -1, 4)

        self.last_point = (x, y)
        self.update_display()

    def mouse_release(self, event):
        if not self.drawing or self.img is None:
            return
        self.drawing = False

        if self.lb_down:
            self.lb_down = False
            self.firt_choose = False
        elif self.rb_down:
            self.rb_down = False

        # 执行 GrabCut
        self.apply_grabcut()

    def apply_grabcut(self):
        bgdModel = np.zeros((1, 65), np.float64)
        fgdModel = np.zeros((1, 65), np.float64)
        rect = (1, 1, self.img.shape[1], self.img.shape[0])
        self.img_gc = self.img.copy()
        cv2.grabCut(self.img_gc, self.mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)
        mask2 = np.where((self.mask == 2) | (self.mask == 0), 0, 1).astype('uint8')
        self.img_gc = self.img_gc * mask2[:, :, np.newaxis]
        self.result_ready = True
        self.update_display(self.img_gc)

    def update_display(self, image=None):
        if image is None:
            image = self.img_show
        if image is None:
            return
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb.shape
        bytes_per_line = ch * w
        qimg = QImage(rgb.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qimg).scaled(
            self.image_label.width() - 20, self.image_label.height() - 20,
            Qt.KeepAspectRatio, Qt.SmoothTransformation
        )
        self.image_label.setPixmap(pixmap)

    def map_to_image(self, x, y):
        """将 QLabel 坐标映射回原始图像坐标"""
        if self.img is None:
            return -1, -1
        label_w, label_h = self.image_label.width() - 20, self.image_label.height() - 20
        img_h, img_w = self.img.shape[:2]
        scale = min(label_w / img_w, label_h / img_h)
        new_w, new_h = int(img_w * scale), int(img_h * scale)
        offset_x = (self.image_label.width() - new_w) // 2
        offset_y = (self.image_label.height() - new_h) // 2
        if not (offset_x <= x <= offset_x + new_w and offset_y <= y <= offset_y + new_h):
            return -1, -1
        img_x = int((x - offset_x) / scale)
        img_y = int((y - offset_y) / scale)
        return img_x, img_y


if __name__ == "__main__":
    app = QApplication(sys.argv)
    # 可以设置应用的全局字体
    font = app.font()
    font.setFamily("SimHei")
    app.setFont(font)

    window = GrabCutApp()
    window.show()
    sys.exit(app.exec_())
